/**
    Copyright (c) 2010-2011 yakiimo02
    Distributed under the New BSD License.
    See included license.txt or http://www.yakiimo3d.com/NewBSDLicense.txt
**/

#include "DXUT.h"
#include "DXHelper.h"
#include "MandelbrotCuda.h"
#include "CudaHelper.h"
#include "DataBufferDX11.h"

#pragma warning(push)
#pragma warning(disable:4201 4408)
#include <cuda.h>
#include <builtin_types.h>
#include <cuda_runtime_api.h>
#include <cuda_d3d11_interop.h>
#pragma warning(pop)

// The CUDA kernel launchers that get called
extern "C" 
{
    void cuda_mandelbrot( void* surface, int width, int height, int pitch );
}

MandelbrotCuda::MandelbrotCuda() : m_nWidth( 0 ), m_nHeight( 0 ), m_pCudaResourceTarget( NULL )
{
}

MandelbrotCuda::~MandelbrotCuda()
{
}

HRESULT MandelbrotCuda::OnCreateDevice( ID3D11Device* pd3dDevice )
{
    HRESULT hr = S_OK;
    return hr;
}

HRESULT MandelbrotCuda::CompileShaders( ID3D11Device* pd3dDevice )
{
    HRESULT hr = S_OK;
    return hr;
}

HRESULT MandelbrotCuda::OnResizedSwapChain( ID3D11Device* pd3dDevice, DataBufferDX11* pDataBuffer )
{
    HRESULT hr = S_OK;

    ID3D11Buffer* pBufferDX11 = pDataBuffer->GetpBufferCUDA();

    m_nWidth = pDataBuffer->GetWidth();
    m_nHeight = pDataBuffer->GetHeight();

    // register the Direct3D resource
    cudaGraphicsD3D11RegisterResource( &m_pCudaResourceTarget, pBufferDX11, cudaGraphicsRegisterFlagsNone );
    CudaCheckForErrors( L"MandelbrotCuda::OnResizedSwapChain cudaGraphicsD3D11RegisterResource failed" );

    //cudaGraphicsResourceSetMapFlags( m_pCudaResourceTarget, cudaGraphicsMapFlagsWriteDiscard );
    //CudaCheckForErrors(L"MandelbrotCuda::OnFrameRender cudaGraphicsResourceSetMapFlags failed");

    return hr;
}

void MandelbrotCuda::OnFrameRender( ID3D11Device* pd3dDevice, ID3D11DeviceContext* pd3dImmediateContext,
                                  double fTime, float fElapsedTime, void* pUserContext )
{
    UpdateConstantBuf( m_nWidth, m_nHeight );

    cudaMemcpyToSymbol( "cb", &m_cb, sizeof(m_cb) );
    CudaCheckForErrors(L"MandelbrotCuda::OnFrameRender cudaMemcpyToSymbol failed");

    cudaStream_t stream = 0;
    const int nResources = 1;
    cudaGraphicsMapResources( nResources, &m_pCudaResourceTarget, stream );
    CudaCheckForErrors(L"MandelbrotCuda::OnFrameRender cudaGraphicsMapResources failed");

    BufType* devPtr;
    size_t size;
    cudaGraphicsResourceGetMappedPointer( reinterpret_cast< void** >( &devPtr ), &size, m_pCudaResourceTarget );
    CudaCheckForErrors(L"MandelbrotCuda::OnFrameRender cudaGraphicsResourceGetMappedPointer failed");

    int nPitch = m_nWidth*sizeof(BufType);

    // kick off the kernel
    cuda_mandelbrot( reinterpret_cast< void* >( devPtr ), m_nWidth, m_nHeight, nPitch );
    CudaCheckForErrors(L"MandelbrotCuda::OnFrameRender cuda_mandelbrot failed");

    cudaGraphicsUnmapResources( nResources, &m_pCudaResourceTarget, stream );
    CudaCheckForErrors(L"MandelbrotCuda::OnFrameRender cudaGraphicsUnmapResources failed");
}

void MandelbrotCuda::OnDestroyDevice()
{
    // unregister the Cuda resources
    cudaGraphicsUnregisterResource( m_pCudaResourceTarget );
    CudaCheckForErrors( L"MandelbrotCuda::OnDestroyDevice() cudaGraphicsUnregisterResource failed" );
}